from flask import Flask, request, Response
from flask_cors import CORS
from openai import OpenAI
import json

app = Flask(__name__)
CORS(app)

# 初始化客户端
client = OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed")

# 提瓦特助手的系统提示词
SYSTEM_PROMPT = """你是一位来自提瓦特大陆的向导派蒙，精通原神游戏的各个方面。
你可以：
1. 提供角色攻略和培养建议
2. 推荐最佳配队搭配
3. 分享深境螺旋攻略
4. 解答活动玩法问题
5. 提供圣遗物搭配建议

请用小派蒙的口吻与旅行者交流。直接输出文本内容，写完一段就换行，不要使用markdown格式，不要使用markdown格式，不要使用markdown格式。"""

@app.route('/chat', methods=['POST'])
def chat_stream():
    # 获取客户端传递的对话历史
    request_data = request.get_json()
    messages = request_data.get('messages', [])
    
    # 如果是新对话，添加系统提示
    if not messages:
        messages.append({
            "role": "system",
            "content": SYSTEM_PROMPT
        })

    # 创建流式对话补全
    completion = client.chat.completions.create(
        model="local-model",
        messages=messages,
        temperature=0.7,
        stream=True
    )

    def event_stream():
        new_message = {"role": "assistant", "content": ""}
        # 流式处理每个数据块
        for chunk in completion:
            if chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                new_message["content"] += content
                yield f"data: {json.dumps({'token': content})}\n\n"
    
    return Response(event_stream(), mimetype='text/event-stream')

if __name__ == '__main__':
    app.run(port=8001)